import torch

""" 模拟单通道二维卷积 """
img = torch.tensor([
    [3, 3, 2, 1, 0],
    [0, 0, 1, 3, 1],
    [3, 1, 2, 2, 3],
    [2, 0, 0, 2, 2],
    [2, 0, 0, 0, 1]
])

kernel = torch.tensor([
    [0, 1, 2],
    [2, 2, 0],
    [0, 1, 2]
])

# 第一次卷积
i = img[0:3, 0:3]
t = (i * kernel).sum()
print(t)

# 模拟整体的卷积
H, W = img.shape
KH, KW = kernel.shape
conv_img = torch.zeros((H - KH + 1, W - KW + 1))
for i in range(H - KH + 1):
    for j in range(W - KW + 1):
        i_img = img[i:i + KH, j:j + KW]
        t = (i_img * kernel).sum()
        conv_img[i, j] = t
print(conv_img)